package com.kongxiang.cubebit.ui.api;

import com.kongxiang.cubebit.base.common.dto.GenericRequest;
import com.kongxiang.cubebit.base.common.dto.R;
import com.kongxiang.cubebit.llm.dto.ReRankRequest;
import com.kongxiang.cubebit.llm.dto.ReRankResponse;
import com.kongxiang.cubebit.llm.service.OpenAIService;
import com.kongxiang.cubebit.llm.service.RagService;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.media.Content;
import io.swagger.v3.oas.annotations.media.Schema;
import io.swagger.v3.oas.annotations.responses.ApiResponse;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.document.Document;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;

@RestController
@RequestMapping("/api/chat")
@Tag(name = "Chat API", description = "提供与聊天相关的API接口")
public class ChatApiController {

    @Autowired
    private OpenAIService openAIService;
    @Autowired
    private RagService ragService;

    @PostMapping("/v1/chat/completions")
    @Schema(description = "根据输入的提示生成聊天响应流")
    public Flux<String> generateChatResponse(@Parameter(description = "聊天提示信息", required = true) @RequestBody GenericRequest<String> prompt) {
        return openAIService.openAiStream(prompt.getBody());
    }

    @PostMapping("/v2/chat/completions")
    @Schema(description = "根据输入的提示生成聊天响应流")
    public Flux<String> chatStream(@Parameter(description = "聊天提示信息", required = true) @RequestBody GenericRequest<String> prompt) {
        return openAIService.openAiStream(prompt.getBody());
    }

    @PostMapping("/v1/embeddings")
    @Operation(summary = "生成嵌入向量", description = "根据输入的提示生成嵌入向量", responses = {
            @ApiResponse(responseCode = "200", description = "成功生成嵌入向量", content = @Content(schema = @Schema(implementation = float[].class))),
            @ApiResponse(responseCode = "400", description = "无效的请求参数")
    })
    public R<float[]> embedding(@Parameter(description = "嵌入提示信息", required = true) @RequestBody GenericRequest<String> prompt) {
        return R.success(openAIService.embedding(prompt.getBody()));
    }

    @PostMapping("/v1/rerank")
    @Operation(summary = "重排序", description = "根据查询条件返回最相似的K个向量结果", responses = {
            @ApiResponse(responseCode = "200", description = "成功重排序", content = @Content(schema = @Schema(implementation = String.class))),
            @ApiResponse(responseCode = "400", description = "无效的请求参数")
    })
    public R<ReRankResponse> rerank(@RequestBody @Validated ReRankRequest request) {
        return R.success(openAIService.rerank(request));
    }

    @PostMapping("/v1/answer")
    @Operation(summary = "问答", description = "根据输入的提示生成问答", responses = {
            @ApiResponse(responseCode = "200", description = "成功生成问答", content = @Content(schema = @Schema(implementation = String.class))),
            @ApiResponse(responseCode = "400", description = "无效的请求参数")
    })
    public R<List<Document>> answer(@Parameter(description = "问答提示信息", required = true) @RequestBody GenericRequest<String> prompt) {
        return R.success(ragService.ragSearch(prompt.getBody(), 5, 0.3));
    }

}